from fastapi import FastAPI, Request, Response
from fastapi.responses import JSONResponse
from typing import Callable
import json

app = FastAPI()

class UnifiedResponseMiddleware:
    def __init__(self, app: FastAPI):
        self.app = app

    async def __call__(self, scope, receive, send):
        async def _send_wrapper(response: Response):
            if response.status_code == 200:
                # 将原始响应内容转换为统一格式
                content = await response.json()
                unified_response = {
                    "code": 200,
                    "msg": "Success",
                    "data": content
                }
                response = JSONResponse(content=unified_response, status_code=200)
            await send(response)

        await self.app(scope, receive, _send_wrapper)

@app.middleware("http")
async def set_fixed_response_structure(request: Request, call_next):
    try:
        # 调用下一个中间件或路由处理函数
        response = await call_next(request)
        if response.status_code == 200:
            # 获取原始响应数据
            if isinstance(response, JSONResponse):
                data = response.body.decode()
                data = json.loads(data)
                print(123)
            else:
                data = await response.body_iterator.__anext__()
                data = data.decode()
                data = json.loads(data)
                print(456)

            # 构建固定结构的响应
            fixed_response = {
                "code": 200,
                "msg": "success",
                "data": data
            }
            return JSONResponse(content=fixed_response)
        return response
    except Exception as e:
        # 发生异常时的固定结构响应
        error_response = {
            "code": 500,
            "msg": str(e),
            "data": None
        }
        return JSONResponse(content=error_response, status_code=500)

# 示例路由
@app.get("/")
async def read_root():
    return {"message": "Hello World"}

@app.get("/items/{item_id}")
async def read_item(item_id: int):
    return {"item_id": item_id}

# 添加中间件
# app.add_middleware(UnifiedResponseMiddleware)